import torch
import numpy as np
import torch.optim as optim


# 张量的创建函数
def create_tensor():
    # print(torch.__version__)  # 应输出版本号（如2.2.1）
    data = [[1, 2], [3, 4]]

    # 直接来自数据
    x_data = torch.tensor(data)
    print(x_data)
    print(type(x_data))

    # 从 NumPy 数组
    np_arr = np.array(data)
    x_np = torch.from_numpy(np_arr)
    print(x_np)

    # 从另一个张量
    # 新张量保留参数张量的属性（形状、数据类型），除非显式覆盖。
    x_ones = torch.ones_like(x_data)  # retains the properties of x_data
    print(f"Ones Tensor: \n {x_ones} \n")

    x_rand = torch.rand_like(x_data, dtype=torch.float)  # overrides the datatype of x_data
    print(f"Random Tensor: \n {x_rand} \n")

    # 使用随机值或常量值
    # shape是张量维度的元组。在下面的函数中，它确定输出张量的维数。
    shape = (2, 3,)
    rand_tensor = torch.rand(shape)
    ones_tensor = torch.ones(shape)
    zeros_tensor = torch.zeros(shape)
    print(f"Random Tensor: \n {rand_tensor} \n")
    print(f"Ones Tensor: \n {ones_tensor} \n")
    print(f"Zeros Tensor: \n {zeros_tensor}")

    print(rand_tensor.device)


# 张量的操作
def tensor_operation():
    # # 这样写是列表
    # e = [[1, 2], [3, 4], [5, 6]]
    # f = [[7, 8], [9, 10], [11, 12]]

    # 1. 创建张量
    e = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.float32)
    f = torch.tensor([[7, 8], [9, 10], [11, 12]], dtype=torch.float32)

    # 张量相加
    print(f"张量相加：{e + f}")

    # 逐元素乘法
    print(f"张量逐元素乘法：{e * f}")

    # 张量的转置
    g = torch.tensor([[1, 2, 3], [4, 5, 6]])
    print(f"张量转置：{g.t()}")  # 或者 g.transpose(0, 1)

    # 张量的形状
    print(f"张量的形状：{g.shape}")  # 返回形状


# pytorch计算梯度
def pytorch_gradient():
    # 创建一个张量，并启用梯度跟踪（requires_grad=True） 2.0表示创建了一个标量（单个数值）张量，其值为浮点数 2.0。
    x = torch.tensor(2.0, requires_grad=True)

    # 定义一个计算
    y = x ** 2 + 3 * x + 1  # y = x² + 3x + 1

    # 计算梯度（反向传播）
    y.backward()

    # 查看x的梯度（dy/dx = 2x + 3）
    print(x.grad)  # 输出：tensor(7.) （因为 2 * 2 + 3 = 7）
    print(type(x.grad))


# 训练一个简单线性回归模型
def train_linear_regression():
    # 数据
    X = torch.tensor([[1.0], [2.0], [3.0]])
    y = torch.tensor([[2.0], [4.0], [6.0]])

    # 模型（线性回归 y = w * x + b）
    w = torch.tensor(0.0, requires_grad=True)
    b = torch.tensor(0.0, requires_grad=True)

    # 优化器（如SGD）
    optimizer = optim.SGD([w, b], lr=0.01)

    # 训练循环
    for epoch in range(100):
        # 前向传播
        y_pred = w * X + b
        loss = ((y_pred - y) ** 2).mean()  # MSE Loss

        # 反向传播
        loss.backward()

        # 更新参数（梯度下降）
        optimizer.step()

        # 清零梯度（重要！）
        optimizer.zero_grad()

    print(f"训练后的参数: w={w.item()}, b={b.item()}")


# 张量的属性
def tensor_property():
    # 创建一个 2D 张量
    tensor = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)

    # 张量的属性
    print("Tensor:\n", tensor)
    print("Shape:", tensor.shape)  # 获取形状
    print("Size:", tensor.size())  # 获取形状（另一种方法）
    print("Data Type:", tensor.dtype)  # 数据类型
    print("Device:", tensor.device)  # 设备
    print("Dimensions:", tensor.dim())  # 维度数
    print("Total Elements:", tensor.numel())  # 元素总数
    print("Requires Grad:", tensor.requires_grad)  # 是否启用梯度
    print("Is CUDA:", tensor.is_cuda)  # 是否在 GPU 上
    print("Is Contiguous:", tensor.is_contiguous())  # 是否连续存储

    # 获取单元素值
    single_value = torch.tensor(42)
    print("Single Element Value:", single_value.item())

    # 转置张量
    tensor_t = tensor.T
    print("Transposed Tensor:\n", tensor_t)


# 张量的基础操作和形状操作
def tensor_basic_operation_and_shape_operation():
    # 创建一个 2D 张量
    tensor = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.float32)
    print("原始张量:\n", tensor)

    # 1. **索引和切片操作**
    print("\n【索引和切片】")
    print("获取第一行:", tensor[0])  # 获取第一行
    print("获取第一行第一列的元素:", tensor[0, 0])  # 获取特定元素
    print("获取第二列的所有元素:", tensor[:, 1])  # 获取第二列所有元素

    # 2. **形状变换操作**
    print("\n【形状变换】")
    reshaped = tensor.view(3, 2)  # 改变张量形状为 3x2
    print("改变形状后的张量:\n", reshaped)
    flattened = tensor.flatten()  # 将张量展平成一维
    print("展平后的张量:\n", flattened)

    # 3. **数学运算操作**
    print("\n【数学运算】")
    tensor_add = tensor + 10  # 张量加法
    print("张量加 10:\n", tensor_add)
    tensor_mul = tensor * 2  # 张量乘法
    print("张量乘 2:\n", tensor_mul)
    tensor_sum = tensor.sum()  # 计算所有元素的和
    print("张量元素的和:", tensor_sum.item())

    # 4. **与其他张量的操作**
    print("\n【与其他张量操作】")
    tensor2 = torch.tensor([[1, 1, 1], [1, 1, 1]], dtype=torch.float32)
    print("另一个张量:\n", tensor2)
    tensor_dot = torch.matmul(tensor, tensor2.T)  # 张量矩阵乘法
    print("矩阵乘法结果:\n", tensor_dot)

    # 5. **条件判断和筛选**
    print("\n【条件判断和筛选】")
    mask = tensor > 3  # 创建一个布尔掩码
    print("大于 3 的元素的布尔掩码:\n", mask)
    filtered_tensor = tensor[tensor > 3]  # 筛选出符合条件的元素
    print("大于 3 的元素:\n", filtered_tensor)


# 主函数
if __name__ == '__main__':
    # create_tensor()
    # tensor_operation()
    # pytorch_gradient()
    # train_linear_regression()
    # tensor_property()
    tensor_basic_operation_and_shape_operation()